// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <complex>
#include <cantor/parser/ast/base.h>
#include <cantor/parser/ast/expr.h>
#include <turbo/memory/arena.h>
#include <cantor/proto/parser/ast.pb.h>

namespace cantor::parser {
    ////////////////////////////////////////////////////////////////////////////////
    /// WhenClauseCase - Represents a single WHEN ... THEN clause inside a CASE expression.
    ///
    /// This AST node corresponds to the SQL syntax:
    ///   WHEN <when_expr> THEN <then_expr>
    ///
    /// It is used exclusively inside CASE expressions to capture conditional branches.
    /// For example:
    ///   CASE
    ///     WHEN score > 90 THEN 'A'
    ///     WHEN score > 80 THEN 'B'
    ///     ELSE 'F'
    ///   END
    ///
    /// Each `WhenClauseCase` holds:
    ///   - `when_expr`: the condition expression after `WHEN`
    ///   - `then_expr`: the result expression after `THEN`
    ///
    /// This node inherits from `ClauseBase` with `CLAUSE_WHEN_CASE` type.
    /// It is part of the AST only and does not represent a directly executable node.
    ///
    /// The `is_complex_node()` method propagates complexity information
    /// based on its child expressions.
    ///
    ////////////////////////////////////////////////////////////////////////////////

    struct WhenClauseCase : public DefiniteExpr {
        WhenClauseCase() : DefiniteExpr(cantor::pb::ast::ENT_CLAUSE_WHEN_CASE, pb::ast::ENC_CLAUSE) {
        }

        nlohmann::ordered_json to_json() const override;

        /// Override to indicate complexity if child expressions are complex
        bool is_complex_node() override;

        /// Condition expression after WHEN
        DispatchExpr *when_expr{nullptr};

        /// Result expression after THEN
        DispatchExpr *then_expr{nullptr};
    };

    ////////////////////////////////////////////////////////////////////////////////
    /// WhenClause - Represents a full CASE expression in the AST.
    ///
    /// This node corresponds to the SQL CASE syntax, supporting both forms:
    ///   1. Simple CASE:
    ///        CASE case_expr
    ///          WHEN val1 THEN result1
    ///          WHEN val2 THEN result2
    ///          ELSE result_default
    ///        END
    ///
    ///   2. Searched CASE:
    ///        CASE
    ///          WHEN cond1 THEN result1
    ///          WHEN cond2 THEN result2
    ///          ELSE result_default
    ///        END
    ///
    /// Fields:
    ///   - `case_expr`:
    ///       Optional expression used in simple CASE form (can be nullptr in searched CASE).
    ///   - `when_clauses`:
    ///       A list of `WhenClauseCase` representing WHEN ... THEN branches.
    ///   - `else_expr`:
    ///       Optional expression representing the ELSE clause (can be nullptr).
    ///
    /// Complexity:
    ///   `is_complex_node()` propagates the complexity flag from its expressions and clauses.
    ///
    /// Usage:
    ///   This node is used during parsing and semantic analysis to build a structured
    ///   representation of CASE expressions. It inherits from `DefiniteExpr`
    ///   with type `CLAUSE_WHEN`.
    ///
    ////////////////////////////////////////////////////////////////////////////////

    struct WhenClause : public DefiniteExpr {
        WhenClause() : DefiniteExpr(pb::ast::ENT_CLAUSE_WHEN_CASE, pb::ast::ENC_CLAUSE) {
        }

        nlohmann::ordered_json to_json() const override;

        /// Check whether any child node is complex
        bool is_complex_node() override;

        /// Optional expression used in simple CASE
        DispatchExpr *case_expr{nullptr};

        /// List of WHEN ... THEN clauses
        Vector<WhenClauseCase *> when_clauses;

        /// Optional ELSE expression
        DispatchExpr *else_expr{nullptr};
    };

    struct TypeStmt;

    struct NamedType : public DefiniteExpr {
        NamedType() : DefiniteExpr(cantor::pb::ast::ENT_NAMED_TYPE, pb::ast::ENC_CLAUSE) {
            name.set_null();
        }

        nlohmann::ordered_json to_json() const override;

        String name;
        TypeStmt *type{nullptr};
    };

    ////////////////////////////////////////////////////////////////////////////////
    /// ConstructArg - Base class for constructor-style arguments in expressions.
    ///
    /// This node is used in AST to represent arguments to complex constructs,
    /// such as struct/map construction, function-like object creation, or
    /// other domain-specific structured expressions.
    ///
    /// It inherits from `ClauseBase` with type `CLAUSE_CONSTRUCT_ARG`.
    ///
    /// Fields:
    ///   - `construct_arg_type`: Defines the specific kind of argument this node represents:
    ///       - `CAT_NONE`        : Undefined / uninitialized.
    ///       - `CAT_MAP_ARG`     : A key-value pair used in map-like construction.
    ///       - `CAT_NAMED_EXPR`  : A named expression, similar to `key = value`.
    ///
    /// Notes:
    ///   - This class is abstract. Specific argument types should inherit and
    ///     specialize further (e.g. `NamedExpr`, `MapPair`, etc).
    ///
    ///   - These are typically used in expression constructs like:
    ///       STRUCT(key1 := val1, key2 := val2)
    ///       MAP('key1', val1, 'key2', val2)
    ///
    ////////////////////////////////////////////////////////////////////////////////

    struct ConstructArg : public DefiniteExpr {
        ConstructArg(pb::ast::ExprNodeType t) : DefiniteExpr(t, pb::ast::ENC_CTOR) {
        }

        nlohmann::ordered_json to_json_impl() const override;
    };

    ////////////////////////////////////////////////////////////////////////////////
    /// MapPair - Represents a key-value pair used in map construction.
    ///
    /// This node is a specialization of `ConstructArg` with type `CAT_MAP_ARG`,
    /// used when parsing expressions that construct map-like structures,
    /// such as: `MAP('key1', val1, 'key2', val2)`
    ///
    /// Fields:
    ///   - `key`   : Expression representing the map key (usually a string or identifier).
    ///   - `value` : Expression representing the corresponding value.
    ///
    /// Notes:
    ///   - Both `key` and `value` are generic `DispatchExpr*`, allowing flexibility
    ///     in what types of expressions can appear.
    ///   - This node is typically only valid in the context of certain constructors,
    ///     and not meant for general use outside map-like AST contexts.
    ///
    ////////////////////////////////////////////////////////////////////////////////

    struct MapPair final : public ConstructArg {
        MapPair() : ConstructArg(pb::ast::ENT_MAP_PAIR) {
        }

        nlohmann::ordered_json to_json() const override;

        DispatchExpr *key{nullptr};
        DispatchExpr *value{nullptr};
    };

    ////////////////////////////////////////////////////////////////////////////////
    /// DictPair - Represents a named argument in a constructor or function call.
    ///
    /// This node is a specialization of `ConstructArg` with type `CAT_NAMED_EXPR`,
    /// used when parsing expressions that provide explicitly named arguments,
    /// such as: `STRUCT(x => 1, y => 2)` or `FUNC(arg1 => value1)`.
    ///
    /// Fields:
    ///   - `name`  : The name of the argument (e.g., "x", "arg1").
    ///   - `value` : The expression associated with this argument name.
    ///
    /// Notes:
    ///   - This is useful for improving clarity and supporting constructors/functions
    ///     that accept named parameters.
    ///   - Only one name-value pair is represented per node.
    ///
    ////////////////////////////////////////////////////////////////////////////////

    struct DictPair final : public ConstructArg {
        DictPair() : ConstructArg(pb::ast::ENT_DICT_PAIR) {
            name = nullptr;
        }

        nlohmann::ordered_json to_json() const override;

        String name;
        DispatchExpr *value{nullptr};
    };

    struct MapArg final : public ConstructArg {
        MapArg() : ConstructArg(pb::ast::ENT_MAP_ARG) {
        }

        nlohmann::ordered_json to_json() const override;

        Vector<MapPair *> values;
    };

    /// for struct/union/json
    /// union is special that only using 0 or 1
    /// of the values list
    struct DictArg final : public ConstructArg {
        DictArg() : ConstructArg(pb::ast::ENT_DICT_ARG) {
        }

        nlohmann::ordered_json to_json() const override;

        Vector<DictPair *> values;
    };

    /// for set/tuple/json array/array
    struct ExprListArg final : public ConstructArg {
        ExprListArg() : ConstructArg(pb::ast::ENT_EXPR_LIST_ARG) {
        }

        nlohmann::ordered_json to_json() const override;

        Vector<DispatchExpr *> values;
    };

    struct TypeNameUri;
    struct EnumArg final : public ConstructArg {
        EnumArg() : ConstructArg(pb::ast::ENT_ENUM_ARG) {
            field = nullptr;
        }

        nlohmann::ordered_json to_json() const override;

        bool is_field{false};
        TypeNameUri *uri{nullptr};
        String field;
        int value{-1};
    };

    struct ProtoArg final : public ConstructArg {
        ProtoArg() : ConstructArg(pb::ast::ENT_PROTO_ARG) {
            uri.set_null();
        }

        nlohmann::ordered_json to_json() const override;

        String uri;
        Vector<DictPair *> values;
    };

    struct UdfArg final : public ConstructArg {
        UdfArg() : ConstructArg(pb::ast::ENT_UDF_ARG) {
        }

        ~UdfArg() override = default;

        nlohmann::ordered_json to_json() const override;

        void set_arg(MapArg * map);

        void set_arg(DictArg *dict);

        void set_arg(ExprListArg* list);

        TypeNameUri *uri{nullptr};
        ConstructArg *arg{nullptr};
        pb::ast::ExprNodeType arg_type{pb::ast::ENT_NONE};
    };

    struct IndexLockOpt : public DefiniteExpr {
        IndexLockOpt() : DefiniteExpr(pb::ast::ENT_INDEX_LOCK_OPT, pb::ast::ENC_CLAUSE) {
        }

        nlohmann::ordered_json to_json() const override;

        StringConstantExpr *lock_mode{nullptr};
    };

    struct AlgorithmOpt : public DefiniteExpr {
        AlgorithmOpt() : DefiniteExpr(pb::ast::ENT_VAR_ASSIGN, pb::ast::ENC_CLAUSE) {
        }

        nlohmann::ordered_json to_json() const override;

        StringConstantExpr *algo{nullptr};
    };

    //////////////////////////////////////////////////////////////////////////////
    /// VarAssign - Represents a variable assignment statement.
    ///
    /// This struct models an assignment operation for setting a variable's value
    /// in SQL. It supports assigning values to variables within SQL scripts or
    /// expressions. The `key` field represents the variable name, and the `value`
    /// field holds the expression being assigned to the variable.
    ///
    /// Fields:
    /// - `key`: The name of the variable to which a value is being assigned.
    /// - `value`: The expression whose result will be assigned to the variable.
    ///
    /// Example:
    /// SET my_variable = 10;  // Assigns the value 10 to `my_variable`
    ///
    //////////////////////////////////////////////////////////////////////////////

    struct VarAssign : public DefiniteExpr {
        String key;
        DispatchExpr *value = nullptr;

        VarAssign() : DefiniteExpr(cantor::pb::ast::ENT_VAR_ASSIGN, pb::ast::ENC_CLAUSE) {
            key = nullptr;
        }

        nlohmann::ordered_json to_json() const override;
    };

    struct NamespaceOption : public DefiniteExpr {
        NamespaceOption() : DefiniteExpr(cantor::pb::ast::ENT_NAMESPACE_OPT, pb::ast::ENC_CLAUSE) {
            str_value.set_null();
        }

        nlohmann::ordered_json to_json() const override;

        pb::ast::NamespaceOptionType type = pb::ast::NAMESPACE_OPT_NONE;
        String str_value;
        uint64_t uint_value = 0;
    };

    struct TableName : public DefiniteExpr {
        TableName() : DefiniteExpr(cantor::pb::ast::ENT_TABLE_NAME, pb::ast::ENC_CLAUSE) {
            ns_name.set_null();
            db_name.set_null();
            table_name.set_null();
        }

        nlohmann::ordered_json to_json() const override;

        void to_stream(std::ostream &os) const override {
            if (!ns_name.empty()) {
                os << ns_name<<".";
            }
            if (!db_name.empty()) {
                os << db_name<<".";
            }
            os <<table_name;
        }
        String ns_name;
        String db_name;
        String table_name;
    };

    struct ColumnName : public DefiniteExpr {
        ColumnName() : DefiniteExpr(cantor::pb::ast::ENT_COLUM_NAME, pb::ast::ENC_CLAUSE) {
            ns_name.set_null();
            db_name.set_null();
            table_name.set_null();
            column_name.set_null();
        }

        nlohmann::ordered_json to_json() const override;

        void to_stream(std::ostream &os) const override {
            if (!ns_name.empty()) {
                os << ns_name<<".";
            }
            if (!db_name.empty()) {
                os << db_name<<".";
            }
            if (!table_name.empty()) {
                os << table_name<<".";
            }
            os <<column_name;
        }
        String ns_name;
        String db_name;
        String table_name;
        String column_name;
    };

    struct TypeNameUri : public DefiniteExpr {
        TypeNameUri() : DefiniteExpr(cantor::pb::ast::ENT_TYPE_NAME_URI, pb::ast::ENC_CLAUSE) {
            ns_name.set_null();
            db_name.set_null();
            table_name.set_null();
            type_name.set_null();
        }

        nlohmann::ordered_json to_json() const override;

        void to_stream(std::ostream &os) const override {
            if (!ns_name.empty()) {
                os << ns_name<<".";
            }
            if (!db_name.empty()) {
                os << db_name<<".";
            }
            if (!table_name.empty()) {
                os << table_name<<".";
            }
            os <<type_name;
        }
        String ns_name;
        String db_name;
        String table_name;
        String type_name;
    };
} // namespace cantor::parser
